"""
矢量化会议调度环境

支持多进程并行执行，提高训练效率
在Windows环境下使用multiprocessing或Ray
"""

import numpy as np
import json
import random
from typing import List, Dict, Tuple, Optional
from multiprocessing import Process, Queue, Event
import sys
import os

# 添加路径以便导入
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))))

from .meeting_agent_env import MeetingSchedulingEnv
from .projection import meeting_projection


class MeetingSchedulingWorker:
    """单个环境Worker"""

    def __init__(self, worker_id: int, seed: int, max_steps: int = 20):
        """
        参数:
            worker_id: Worker ID
            seed: 随机种子
            max_steps: 最大步数
        """
        self.worker_id = worker_id
        self.seed = seed + worker_id
        self.max_steps = max_steps

        # 创建环境
        self.env = MeetingSchedulingEnv(
            max_steps=max_steps,
            num_total_people=6,
            seed=self.seed
        )

        self.current_observation = None
        self.current_info = {}

    def reset(self, instance: Optional[Dict] = None):
        """重置环境"""
        observation, info = self.env.reset(instance)
        self.current_observation = observation
        self.current_info = info
        return observation, info

    def step(self, action_dict: Dict, valid: int):
        """执行步骤"""
        observation, reward, done, info = self.env.step(action_dict, valid)
        self.current_observation = observation
        self.current_info = info
        return observation, reward, done, info

    def get_state(self):
        """获取当前状态"""
        return self.env.get_state_dict()


def worker_process(worker_id: int,
                   seed: int,
                   max_steps: int,
                   command_queue: Queue,
                   result_queue: Queue,
                   ready_event: Event):
    """
    Worker进程函数

    参数:
        worker_id: Worker ID
        seed: 随机种子
        max_steps: 最大步数
        command_queue: 命令队列（接收主进程指令）
        result_queue: 结果队列（返回结果给主进程）
        ready_event: 就绪事件
    """
    # 创建Worker
    worker = MeetingSchedulingWorker(worker_id, seed, max_steps)

    # 标记就绪
    ready_event.set()

    # 主循环
    while True:
        try:
            command, data = command_queue.get()

            if command == 'reset':
                instance = data.get('instance')
                obs, info = worker.reset(instance)
                result_queue.put(('reset', {'obs': obs, 'info': info}))

            elif command == 'step':
                action_dict = data.get('action_dict')
                valid = data.get('valid')
                obs, reward, done, info = worker.step(action_dict, valid)
                result_queue.put(('step', {
                    'obs': obs,
                    'reward': reward,
                    'done': done,
                    'info': info
                }))

            elif command == 'get_state':
                state = worker.get_state()
                result_queue.put(('state', {'state': state}))

            elif command == 'close':
                break

        except Exception as e:
            result_queue.put(('error', {'error': str(e)}))
            break


class MeetingSchedulingMultiProcessEnv:
    """
    多进程矢量化会议调度环境

    支持并行执行多个环境实例，提高采样效率
    """

    def __init__(self,
                 seed: int,
                 env_num: int,
                 max_steps: int = 20,
                 dataset_path: Optional[str] = None,
                 is_train: bool = True):
        """
        参数:
            seed: 随机种子
            env_num: 环境数量
            max_steps: 每个环境的最大步数
            dataset_path: 数据集路径（可选，用于加载预定义实例）
            is_train: 是否为训练模式
        """
        self.seed = seed
        self.env_num = env_num
        self.max_steps = max_steps
        self.dataset_path = dataset_path
        self.is_train = is_train

        # 加载数据集（如果提供）
        self.instances = []
        if dataset_path and os.path.exists(dataset_path):
            with open(dataset_path, 'r', encoding='utf-8') as f:
                all_instances = json.load(f)

            # 划分训练/验证集
            split_idx = int(len(all_instances) * 0.8)
            if is_train:
                self.instances = all_instances[:split_idx]
            else:
                self.instances = all_instances[split_idx:]

        self.data_len = len(self.instances) if self.instances else 0

        # 初始化随机数生成器
        self._rng = np.random.RandomState(seed)

        # 不使用多进程（Windows兼容性），直接创建Worker列表
        self.workers = []
        for i in range(env_num):
            worker = MeetingSchedulingWorker(i, seed, max_steps)
            self.workers.append(worker)

        # 当前状态
        self.current_observations = [None] * env_num
        self.current_infos = [{}] * env_num
        self.dones = [False] * env_num

    def reset(self) -> Tuple[List[str], List[Dict]]:
        """
        重置所有环境

        返回:
            (observations, infos): 观察列表和信息列表
        """
        observations = []
        infos = []

        for i, worker in enumerate(self.workers):
            # 随机选择实例（如果有数据集）
            if self.instances:
                instance = self._rng.choice(self.instances)
            else:
                instance = None

            obs, info = worker.reset(instance)
            observations.append(obs)
            infos.append(info)

            self.current_observations[i] = obs
            self.current_infos[i] = info
            self.dones[i] = False

        return observations, infos

    def step(self,
            actions: List[str],
            valids: List[int]) -> Tuple[List[str], List[float], List[bool], List[Dict]]:
        """
        并行执行所有环境的步骤

        参数:
            actions: 文本动作列表（LLM输出）
            valids: 有效性标志列表

        返回:
            (observations, rewards, dones, infos)
        """
        # 投影文本动作到字典格式
        action_dicts, projected_valids = meeting_projection(actions)

        # 合并valids（如果projection判断无效，则覆盖）
        final_valids = [
            v1 and v2 for v1, v2 in zip(valids, projected_valids)
        ]

        observations = []
        rewards = []
        dones = []
        infos = []

        for i, worker in enumerate(self.workers):
            if self.dones[i]:
                # 已完成的环境保持原状态
                observations.append(self.current_observations[i])
                rewards.append(0.0)
                dones.append(True)
                infos.append(self.current_infos[i])
            else:
                # 执行步骤
                obs, reward, done, info = worker.step(
                    action_dicts[i],
                    final_valids[i]
                )

                observations.append(obs)
                rewards.append(reward)
                dones.append(done)
                infos.append(info)

                self.current_observations[i] = obs
                self.current_infos[i] = info
                self.dones[i] = done

        return observations, rewards, dones, infos

    def get_states(self) -> List[Dict]:
        """获取所有环境的状态"""
        return [worker.get_state() for worker in self.workers]

    def close(self):
        """关闭所有环境"""
        # 在这个简化版本中，不需要特殊关闭操作
        pass

    def __len__(self):
        return self.env_num


def test_multiprocess_env():
    """测试矢量化环境"""
    print("=== 测试矢量化环境 ===\n")

    # 创建环境
    env = MeetingSchedulingMultiProcessEnv(
        seed=42,
        env_num=2,
        max_steps=20
    )

    print(f"环境数量: {len(env)}\n")

    # 重置
    print("=== 重置环境 ===")
    observations, infos = env.reset()

    for i, (obs, info) in enumerate(zip(observations, infos)):
        print(f"\n环境 {i}:")
        print(f"观察: {obs[:100]}...")
        print(f"信息: {info}")

    # 步骤1: 查询日程
    print("\n\n=== 步骤1: 查询日程 ===")

    actions = [
        '''
        首先查询所有参会人员的日程。
        ```json
        {
            "action_type": "query_calendar",
            "parameters": {
                "persons": ["李明", "王芳", "陈静"]
            }
        }
        ```
        ''',
        '''
        查询必选人员的日程。
        ```json
        {
            "action_type": "query_calendar",
            "parameters": {
                "persons": ["李明", "王芳"]
            }
        }
        ```
        '''
    ]

    valids = [1, 1]

    observations, rewards, dones, infos = env.step(actions, valids)

    for i, (obs, reward, done, info) in enumerate(zip(observations, rewards, dones, infos)):
        print(f"\n环境 {i}:")
        print(f"观察: {obs[:150]}...")
        print(f"奖励: {reward}")
        print(f"完成: {done}")

    # 步骤2: 查询历史
    print("\n\n=== 步骤2: 查询历史 ===")

    actions = [
        '''
        ```json
        {
            "action_type": "query_history",
            "parameters": {}
        }
        ```
        ''',
        '''
        ```json
        {
            "action_type": "query_history",
            "parameters": {}
        }
        ```
        '''
    ]

    valids = [1, 1]

    observations, rewards, dones, infos = env.step(actions, valids)

    for i, (obs, reward, done, info) in enumerate(zip(observations, rewards, dones, infos)):
        print(f"\n环境 {i}:")
        print(f"观察: {obs[:150]}...")
        print(f"奖励: {reward}")

    # 关闭环境
    env.close()
    print("\n\n环境已关闭")


if __name__ == '__main__':
    test_multiprocess_env()
